# coding=UTF-8
import os
import time
import uuid
from xml.dom.minidom import Document
import getopt, sys
import itertools
from itertools import combinations, filterfalse
import glob
import glob_jmr
import glob_iot

def usage():
    print("[usage]:")
    print("python test.py [-h] -t [type] -o [output_path]")
    print("[-t]:")
    print("\t0101: JMR恶意资源(0xe8)")
    print("\t0102: JMR恶意报文(0xe9")
    print("\t0103: JMR恶意文件(0xea)")
    print("\t0201: JMR处置事件(0xeb)")
    print("\t01  : IOT恶意报文(0xe9")
    print("\t02  : IOT恶意文件(0xea)")
    print("\t03  : IOT流量筛查(0xed)")

################################################################################
# include depend file 
glob._init()
glob_iot._init()
glob_jmr._init()
def test_list():
    lst=[10,20,30,40,50]
    for i in range(len(lst)):
        print("index=%u: elem=%s" % (i,lst[i]))
    lst=[[10, '10a'], [20, '20a'], [30, '30a'], [40, '40a'] ,[50, '50a'] ]
    for i in range(len(lst)):
        val=lst[i]
        print("index=%u: elem=%u,%s" % (i, val[0], val[1]))
def test_dict():
    dt = [ {'type_01' : '01'}, {'type_02' : '02'}, {'type_03' : '03'}] 
    for i in range(len(dt)):
        for k, v in dt[i].items():
            print("index=%u: %s:%s" % (i, k, v))

################################################################################
# common function
def add_elem(parent_node, node_name, node_value=None):
    """
    为父节点对象添加子节点对象
    :param node_name:  子节点名称
    :param parent_node:      父节点对象
    :param node_value:  子节点内容文本
    :return:
    """
    curr_elem = doc.createElement(node_name)
    parent_node.appendChild(curr_elem)
    if node_value:
        curr_elem.appendChild(doc.createTextNode(node_value))


def add_xml_header(doc, command_node):
    global cmd_info_table
    global cmd_object_table
    global cmd_result_table
    # create xml node: commandInfo
    command_info_node = doc.createElement('commandInfo')
    command_node.appendChild(command_info_node)
    for i in range(len(cmd_info_table)):
        for k, v in cmd_info_table[i].items():
            add_elem(parent_node=command_info_node, node_name=k, node_value=v[0])

    # create xml node: commandObject
    command_object_node = doc.createElement('commandObject')
    command_node.appendChild(command_object_node)
    for i in range(len(cmd_object_table)):
        for k, v in cmd_object_table[i].items():
            add_elem(parent_node=command_object_node, node_name=k, node_value=v[0])

    # create xml node: commandResult
    command_result_node= doc.createElement('commandResult')
    command_node.appendChild(command_result_node)
    for i in range(len(cmd_result_table)):
        for k, v in cmd_result_table[i].items():
            add_elem(parent_node=command_object_node, node_name=k, node_value=v[0])

################################################################################
# main function
param_cmd_type = None
param_output_path = './'
def parse_main_param():
    global param_cmd_type
    global param_output_path
    try:
        opts, args = getopt.getopt(sys.argv[1:], "ht:o:", ["help", "type=", "output="])
    except getopt.GetoptError as err:
        print(err)
        usage()
        sys.exit(2)
    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("-t", "--type"):
            param_cmd_type = a 
        elif o in ("-o", "--output"):
            param_output_path = a 
        else:
            usage()
            sys.exit(2)

if __name__ == "__main__":
    parse_main_param()
    # check param: output path
    if (not param_output_path) or (not os.path.isabs(param_output_path)):
        param_output_path = os.getcwd()
    if not os.path.exists(param_output_path):
        os.makedirs(param_output_path)

    # check param: cmd_type
    if not param_cmd_type:
        print("[ERROR]: input param error")
        usage()
        sys.exit(1)
    command_id = glob.command_id
    cmd_type_table = glob.cmd_type_table
    cmd_type_list=None
    for i in range(len(cmd_type_table)):
        if cmd_type_table[i][0] ==  param_cmd_type:
            cmd_type_list = cmd_type_table[i]
            rule_name = cmd_type_list[2]
            break
    if not cmd_type_list:
        print("[ERROR]: command type error: command_type=%s: unknown command!!!" % (param_cmd_type))
        sys.exit()
    print("[param]: type=%s, output_path=%s" % (param_cmd_type, param_output_path))
    if param_cmd_type == "0101":
        cmd_info_table = glob_jmr.cmd_info_table
        cmd_object_table = glob_jmr.cmd_object_table
        cmd_result_table = glob_jmr.cmd_result_table
        rule_table = glob_jmr.jmr_evil_resource_rule_table
    elif param_cmd_type == "0102":
        cmd_info_table = glob_jmr.cmd_info_table
        cmd_object_table = glob_jmr.cmd_object_table
        cmd_result_table = glob_jmr.cmd_result_table
        rule_table = glob_jmr.jmr_evil_payload_rule_table
    elif param_cmd_type == "0103":
        cmd_info_table = glob_jmr.cmd_info_table
        cmd_object_table = glob_jmr.cmd_object_table
        cmd_result_table = glob_jmr.cmd_result_table
        rule_table = glob_jmr.jmr_evil_file_rule_table
    elif param_cmd_type == "0201":
        cmd_info_table = glob_jmr.cmd_info_table
        cmd_object_table = glob_jmr.cmd_object_table
        cmd_result_table = glob_jmr.cmd_result_table
        rule_table = glob_jmr.jmr_dispose_rule_table
    elif param_cmd_type == "01":
        cmd_info_table = glob_iot.cmd_info_table
        cmd_object_table = glob_iot.cmd_object_table
        cmd_result_table = glob_iot.cmd_result_table
        rule_table = glob_iot.iot_evil_payload_rule_table
    elif param_cmd_type == "02":
        cmd_info_table = glob_iot.cmd_info_table
        cmd_object_table = glob_iot.cmd_object_table
        cmd_result_table = glob_iot.cmd_result_table
        rule_table = glob_iot.iot_evil_file_rule_table
    elif param_cmd_type == "03":
        cmd_info_table = glob_iot.cmd_info_table
        cmd_object_table = glob_iot.cmd_object_table
        cmd_result_table = glob_iot.cmd_result_table
        rule_table = glob_iot.iot_data_select_rule_table
    else:
        print("[ERROR]: command table error: cannot find command table: command_type=%s" % (param_cmd_type))
        sys.exit()

    # make xml from rules
    rule_opt_list=list(itertools.filterfalse(lambda x : x[0]!=0, rule_table))
    for count in range(1, len(rule_opt_list)+1):
        comb = itertools.combinations(rule_opt_list, count)
        for comb_elem in comb:
            command_id =  int(command_id) + 1
            opt_elem_list = list(comb_elem)
            doc = Document()
            # create xml node: command
            command_node = doc.createElement('command')
            doc.appendChild(command_node)

            # create xml node: commandInfo/commandObject/comandResult
            add_xml_header(doc=doc, command_node=command_node)

            # create xml node: commandRule/commandType
            command_rule_node= doc.createElement('commandRule')
            command_node.appendChild(command_rule_node)
            for i in range(len(cmd_type_table)):
                lst = cmd_type_table[i]
                if lst[0] ==  param_cmd_type:
                    add_elem(parent_node=command_rule_node, node_name=lst[1], node_value=lst[0])

            # create xml node: commandRule/rule...
            rule_node= doc.createElement('rule')
            command_rule_node.appendChild(rule_node)
            for i in range(0, len(rule_table)):
                rule_elem = rule_table[i]
                if rule_elem[0] == 1:
                    add_elem(parent_node=rule_node, node_name=rule_elem[1], node_value=rule_elem[2])
                else:
                    if rule_elem in opt_elem_list:
                        add_elem(parent_node=rule_node, node_name=rule_elem[1], node_value=rule_elem[2])
                    else:
                        add_elem(parent_node=rule_node, node_name=rule_elem[1], node_value='')

            #write to file
            f_name_full = f'{param_output_path}/test_{rule_name}_{glob.command_id}.xml'
            with open(f_name_full, 'w', encoding='utf-8') as f:
                doc.writexml(f, indent='', newl='\n', addindent='    ', encoding='utf-8')
            print(f"Write to {f_name_full} OK")
            doc.unlink()
################################################################################
#ending
